Repository navigation
Conversation
DataLoaderStateMixin.begin() used len(dataset) for the final-batch remainder. With SubsetRandomSampler (or any sized sampler) over a larger map-style dataset, that drops or duplicates samples in gather_for_metrics. Prefer batch_sampler.sampler length when available, then fall back to dataset length / total_dataset_length. Fixes huggingface#4410
| getattr(self, "sampler", None), | ||
| ) | ||
| for sampler in sampler_candidates: | ||
| if sampler is None or isinstance(sampler, BatchSampler): |
There was a problem hiding this comment.
With BatchSampler(drop_last=True), DataLoader.drop_last stays False. Selecting 10 samples with batch size 4 correctly produces 8 samples but this change makes gather_for_metrics keep only 6. Could we check the batch sampler’s setting and add a test?
There was a problem hiding this comment.
Good catch — fixed in 49c2569.
begin() now uses _effective_drop_last() which also walks batch_sampler.drop_last (and nested batch_sampler.batch_sampler), so when BatchSampler(drop_last=True) is passed and DataLoader.drop_last stays False, remainder stays -1 and gather_for_metrics keeps all 8 yielded samples.
Added test_remainder_respects_batch_sampler_drop_last covering the 10→8 case.
| with suppress(Exception): | ||
| if not self._drop_last: | ||
| length = getattr(self.dataset, "total_dataset_length", len(self.dataset)) | ||
| length = self._length_for_remainder() |
There was a problem hiding this comment.
begin() runs before set_epoch(), so this can use the previous epoch’s sampler length. I reproduced a sampler growing from 5 to 6 samples where the sixth sample gets dropped from metrics. Could we also set the epoch first and cover this in a test?
There was a problem hiding this comment.
Agreed — fixed in the same commit.
DataLoaderShard.__iter__ and DataLoaderDispatcher.__iter__ now call set_epoch(self.iteration) before begin(), so remainder uses the current epoch’s sampler length.
Added test_remainder_uses_length_after_set_epoch with a growing sampler (len 5 → 6 across iterations); second pass now keeps all 6 samples.
Michael-RDev
left a comment
Michael-RDev
left a comment
There was a problem hiding this comment.
This fixes the reported subset-sampler case but I found two cases where valid samples get dropped from metrics: batch samplers using ``drop_last=True` and samplers whose length changes between epochs.
Both reproductions pass on the parent commit and fail here. I’ve left some details, please address them and add regression tests :)
- Treat BatchSampler(drop_last=True) as effective drop_last so gather_for_metrics does not truncate full batches when DataLoader.drop_last stays False. - Call set_epoch before begin so remainder uses the current epoch's sampler length (covers epoch-dependent/growing samplers). - Add regression tests for both cases.
|
Addressed both review notes from @Michael-RDev in 49c2569:
Local: |
|
Friendly bump for maintainers: both review notes from @Michael-RDev are addressed in 49c2569 ( |
2 participants
Summary
Fixes #4410.
DataLoaderStateMixin.begin()computed the final-batch remainder fromlen(dataset). When a sized sampler (e.g.SubsetRandomSampler) selects fewer indices than the underlying map-style dataset,gather_for_metricsthen drops or duplicates samples.Change: prefer
len(batch_sampler.sampler)(then nested / top-level sampler) when sized; otherwise keeptotal_dataset_length/len(dataset). AfterAccelerator.prepare, the real selection often lives onbatch_sampler.samplerwhileself.samplercan be a defaultSequentialSamplerover the full dataset — so that order matters.Validation
New test reproduces the issue's SubsetRandomSampler case and the Subset control.
Coordination
Comment left on #4410: issue was unassigned with no open PR; happy to close if @gorkadabo1 prefers to take it.